{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import json\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_url = \"http://localhost:8000/\"\n",
    "task_api = base_url + \"task/\"\n",
    "agent_api = base_url + \"agent/\"\n",
    "cluster_api = base_url + \"cluster/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_request(url, params=None):\n",
    "    response = requests.get(url, params=params)\n",
    "    return response.json()\n",
    "\n",
    "def b64encode(data: str, encoding=\"utf-8\") -> str:\n",
    "    return base64.b64encode(data.encode(encoding)).decode(encoding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务提交"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def submit_task(stage: str, base_model_dir: str = None, resume: bool = False) -> str:\n",
    "    with open(f\"../configs/{stage}_example.json\", \"r\") as f:\n",
    "        config = json.load(f)\n",
    "\n",
    "    if \"stop\" in config[\"task_config\"][stage]:\n",
    "        config[\"task_config\"][stage][\"stop\"][\"training_iteration\"] = 100\n",
    "\n",
    "    if base_model_dir is not None:\n",
    "        config[\"task_config\"][stage][\"base_model\"] = {\"checkpoint_dir\": base_model_dir}\n",
    "    else:\n",
    "        config[\"task_config\"][stage].pop(\"base_model\", None)\n",
    "\n",
    "    url = task_api + (\"resume/{}\" if resume else \"submit/{}\")\n",
    "    encode_data = b64encode(json.dumps(config))\n",
    "    job_id = do_request(url.format(encode_data))\n",
    "    return job_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 预训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrain_job_id = submit_task(\"pretrain\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 博弈演进"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_dir = \"/mnt/ray/ray_results/pretrain_example/DQN_LocalPathPlanning_4965b_00000_0_2024-09-10_03-56-09/checkpoint_000004\"\n",
    "job_id = evolve_job_id = submit_task(\"evolve\", ckpt_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 迁移学习"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_dir = \"/mnt/ray/ray_results/pretrain_example/DQN_LocalPathPlanning_4965b_00000_0_2024-09-10_03-56-09/checkpoint_000004\"\n",
    "transfer_job_id = submit_task(\"transfer\", ckpt_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 推理应用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_dir = \"/mnt/ray/ray_results/fe6503f9-602d-11ef-9344-c85ea95f6813/DQN_LocalPathPlanning_fc8b1_00000_0_2024-09-14_05-00-50/checkpoint_000004\"\n",
    "infer_job_id = submit_task(\"inference\", ckpt_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务状态"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "status = do_request(task_api + f\"status/{job_id}\")\n",
    "status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 中止任务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stopped = do_request(task_api + f\"stop/{job_id}\")\n",
    "stopped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 接续任务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resumed_id = submit_task(\"evolve\", resume=True)\n",
    "resumed_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 删除任务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 仅通过JobSubmissionClient提交的任务可以删除\n",
    "# 运行中的任务无法删除\n",
    "deleted = do_request(task_api + f\"delete/{job_id}\")\n",
    "deleted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务日志"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log = do_request(task_api + f\"log/{job_id}\")\n",
    "print(log[-20000:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务进度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "progress = do_request(task_api + f\"progress/{job_id}\")\n",
    "progress"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"../configs/pretrain_example.json\", \"r\") as f:\n",
    "    config = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = do_request(task_api + f\"result/{config['task_id']}\")\n",
    "# result = json.loads(result)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = pd.DataFrame.from_dict(json.loads(result))\n",
    "metric = \"episode_reward_mean\"\n",
    "df.filter([metric, \"checkpoint_dir_name\"]).groupby(\"checkpoint_dir_name\").mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_url = do_request(task_api + f\"visualize/{config['task_id']}\")\n",
    "print(tensorboard_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 列出任务所有检查点"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoints = do_request(task_api + f\"checkpoint/{config['task_id']}\")\n",
    "checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 保存模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_dir = checkpoints[0]\n",
    "checkpoint_dir = b64encode(checkpoint_dir)\n",
    "save_dir = b64encode(f\"/mnt/ray/export_models/{config['task_name']}_{config['task_id']}\")\n",
    "onnx_version = 17\n",
    "\n",
    "url = agent_api + f\"export/{checkpoint_dir}?save_dir={save_dir}\"\n",
    "if onnx_version:\n",
    "    url += f\"&onnx={onnx_version}\"\n",
    "\n",
    "saved = do_request(url)\n",
    "saved"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务列表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_list = do_request(task_api + \"list\")\n",
    "task_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 节点信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_ip = \"172.19.0.3\"\n",
    "node_url = do_request(cluster_api + f\"node/{node_ip}\")\n",
    "node_url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 集群信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_info = do_request(cluster_api + \"info\")\n",
    "cluster_info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 集群资源"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_resources = do_request(cluster_api + \"resources\")\n",
    "cluster_resources"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ray",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
